import eventlet
import struct
import encoding
import socket
from Crypto.Random.random import StrongRandom
PKCS1_OAEP = eventlet.import_patched('Crypto.Cipher.PKCS1_OAEP')
RSA = eventlet.import_patched('Crypto.PublicKey.RSA')
RIPEMD = eventlet.import_patched('Crypto.Hash.RIPEMD')
HMAC = eventlet.import_patched('Crypto.Hash.HMAC')
enigma = eventlet.import_patched('enigma')
EnigmaMachine = enigma.EnigmaMachine

listen_at = ('0.0.0.0', 3967)
edge = ('127.0.0.1', 1080)
RSA_keyfile = '.\privkey.key'
RSA_data_length = 192
GreenPool_size = 10000
enigma_settings = {'count': GreenPool_size, 'rotors': 'II VI V', 'ring_settings': "1 16 7", 'reflector': 'B'}
client_conn_waiting_time = 0.01 # in seconds
add_padding = True

with open(RSA_keyfile, 'r') as RSA_key:
    key = RSA.importKey(RSA_key.read())
    rsa = PKCS1_OAEP.new(key)

CLOSED_EDGE_CONNECTION = 'CLOSED_EDGE_CONNECTION'

random = StrongRandom()

EnigmaMachine.wrapper_initialize(**enigma_settings)

def handle(conn, addr):
    print('Got a connection from ', addr)
    inbuf = []
    outbuf = []
    edgeconn = []
    ib = ''

    seed = struct.pack("!I", random.randint(0, 4294967294))
    padding_count = random.randint(0, 17)
    for i in xrange(padding_count):
        seed += struct.pack("!I", random.getrandbits(32))

    try:
        conn.send(seed)
    except:
        print('Exception sending to ', addr)
        conn.close()
        return

    try:
        d = rsa.decrypt(conn.recv(1024)[0:RSA_data_length])
    except:
        print('Exception decrypting handshake from ', addr)
        conn.close()
        return

    if len(d) < 35:
        print('Got bad handshake data from ', addr)
        conn.close()
        return

    if not d[0:4] == seed[0:4]:
        print('Got bad handshake data from ', addr)
        conn.close()
        return

    enigmakey = d[4:7]
    alphabet = d[7:31]
    hmackey = d[31:35]

    upstream_hmac = HMAC.new(hmackey, digestmod=RIPEMD.new()) # used to sign data waiting to be sent
    downstream_hmac = HMAC.new(hmackey, digestmod=RIPEMD.new()) # used to verify data we have received

    try:
        upstream_enigma = EnigmaMachine() # used to encrypt data waiting to be sent
        upstream_enigma.set_display(enigmakey)
        downstream_enigma = EnigmaMachine() # used to decrypt data we have received
        downstream_enigma.set_display(enigmakey)
    except:
        try:
            upstream_enigma.destroy()
            downstream_enigma.destroy()
        except:
            pass
        print('Exception initializing enigma machine of ', addr)
        conn.close()
        return

    conn.settimeout(client_conn_waiting_time)

    padding = ''
    padding_count = random.randint(0, 8)
    for i in xrange(padding_count):
        padding += struct.pack("!I", random.getrandbits(32))

    vf = list(upstream_enigma.process_text(alphabet + padding).lower())

    for i in [random.randint(0, len(vf)-1) for i in xrange(random.randint(0, 8))]:
        vf[i] = vf[i].upper()

    [vf.insert(random.randint(1, len(vf)-1), ' ') for i in xrange(random.randint(0, 8))]

    try:
        conn.send(''.join(vf))
    except:
        upstream_enigma.destroy()
        downstream_enigma.destroy()
        print('Exception sending to ', addr)
        conn.close()
        return

    while True:
        try:
            data = conn.recv(65536)
            if not data:
                break
            data = ib + downstream_enigma.process_text(data.replace(' ', ''))
            ib = ''
            if not len(data) % 2 == 0:
                data, ib = data[0:-1], data[-1]
            inbuf.append(encoding.debase(data, alphabet))
        except socket.error:
            pass
        except:
            upstream_enigma.destroy()
            downstream_enigma.destroy()
            print('Exception decoding data of ', addr)
            conn.close()
            return

        if not inbuf == []:
            idata = ''.join(inbuf)
            inbuf = []

            while True:
                if not idata:
                    break

                try:
                    length, = struct.unpack("!I", idata[1:5])
                except:
                    upstream_enigma.destroy()
                    downstream_enigma.destroy()
                    print('Exception decoding data of ', addr)
                    conn.close()
                    return
                if len(idata)-5 < length:
                    inbuf.append(idata)
                    break

                msg, idata = idata[5:5+length], idata[5+length:]

                hmac = msg[0:20]
                downstream_hmac.update(msg[20:])

                if not downstream_hmac.digest() == hmac:
                    upstream_enigma.destroy()
                    downstream_enigma.destroy()
                    print('Got a message with bad HMAC from ', addr)
                    conn.close()
                    return

                if msg[20] == 'c': # create a connection
                    l = len(edgeconn)
                    if l > 4294967295:
                        print('Edge connection id overflow.(4294967296 connections made?)')
                    else:
                        try:
                            sock = eventlet.connect(edge)
                        except:
                            print('Exception connecting to server.')
                            continue
                        edgeconn.append(sock)
                        sock.setblocking(0)
                        outbuf.append('c' + struct.pack("!I", l))
                elif msg[20] == 'd': # close a connection
                    try:
                        conn_id, = struct.unpack("!I", msg[21:25])
                    except:
                        upstream_enigma.destroy()
                        downstream_enigma.destroy()
                        print('Exception decoding data of ', addr, ' while processing close_connection.')
                        conn.close()
                        return

                    if edgeconn[conn_id] == CLOSED_EDGE_CONNECTION:
                        print('Client trying closing a closed connection ', addr)
                        continue
                    else:
                        edgeconn[conn_id].close()
                        edgeconn[conn_id] = CLOSED_EDGE_CONNECTION
                elif msg[20] == 'm': # send to edge
                    try:
                        conn_id, = struct.unpack("!I", msg[21:25])
                    except:
                        upstream_enigma.destroy()
                        downstream_enigma.destroy()
                        print('Exception decoding data of ', addr, ' while processing send_server.')
                        conn.close()
                        return

                    if edgeconn[conn_id] == CLOSED_EDGE_CONNECTION:
                        print('Client trying sending to a closed connection ', addr)
                        continue
                    else:
                        try:
                            edgeconn[conn_id].sendall(msg[25:])
                        except:
                            edgeconn[conn_id] = CLOSED_EDGE_CONNECTION
                            outbuf.append('d' + msg[21:25])
                elif msg[20] == 'p':
                    pass
                else:
                    return

        for i in xrange(len(edgeconn)): # manually select()
            if edgeconn[i] == CLOSED_EDGE_CONNECTION:
                continue
            try:
                data = edgeconn[i].recv(24000)
            except:
                pass
            else:
                if not data:
                    edgeconn[i] = CLOSED_EDGE_CONNECTION
                    outbuf.append('d' + struct.pack("!I", i))
                else:
                    outbuf.append('r' + struct.pack("!I", i) + data)

        if add_padding:
            if random.randint(0, 42) < 3:
                padding_count = random.randint(0, 42)
                padding = ''
                for i in xrange(padding_count):
                    padding += struct.pack("!I", random.getrandbits(32))
                outbuf.append('p' + padding)

        ob = []
        for msg in outbuf:
            upstream_hmac.update(msg)
            hmac = upstream_hmac.digest()

            length = struct.pack("!I", len(hmac) + len(msg))

            encdata = encoding.add_space(enigma.add_uppercase(upstream_enigma.process_text(encoding.enbase(random.choice(alphabet) + length + hmac + msg, alphabet))))

            ob.append(encdata)

        if ob == []:
            continue
        else:
            try:
                conn.sendall(''.join(ob))
            except:
                upstream_enigma.destroy()
                downstream_enigma.destroy()
                print('Exception sending to ', addr)
                conn.close()
                return

        outbuf = []

if __name__ == '__main__':
    eventlet.monkey_patch()
    eventlet.serve(eventlet.listen(listen_at), handle, GreenPool_size)